module spectra

use, intrinsic :: iso_c_binding
use basic
use avd
use response
implicit none

contains

subroutine spamixed(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    integer :: m

    m = 1
    do while ( P(m)<=MPR*dt )
        m = m + 1
    end do

    call spafreq(acc,n,dt,zeta,P(1:m),m,SPA(1:m),SPI(1:m))
    if ( m<nP ) then
        call spanmk(acc,n,dt,zeta,P(m+1:nP),(nP-m),SPA(m+1:nP),SPI(m+1:nP))
    end if

end subroutine spamixed

subroutine pspamixed(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    integer :: m

    m = 1

    do while ( P(m)<=MPR*dt )
        m = m + 1
    end do

    call pspafreq(acc,n,dt,zeta,P(1:m),m,SPA(1:m),SPI(1:m))
    if ( m<nP ) then
        call pspanmk(acc,n,dt,zeta,P(m+1:nP),(nP-m),SPA(m+1:nP),SPI(m+1:nP))
    end if

end subroutine pspamixed

subroutine spavdmixed(acc,n,dt,zeta,P,nP,SPA,SPI,SPV,SPD,SPEV) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n), SPV(n), SPD(n), SPEV(n)
    integer, intent(out) :: SPI(n)

    integer :: m
    m = 1
    do while ( P(m)<=MPR*dt )
        m = m + 1
    end do

    call spavdfreq(acc,n,dt,zeta,P(1:m),m,SPA(1:m),SPI(1:m),&
        SPV(1:m),SPD(1:m),SPEV(1:m))
    if ( m<nP ) then
        call spavdnmk(acc,n,dt,zeta,P(m+1:nP),(nP-m),SPA(m+1:nP),&
            SPI(m+1:nP),SPV(m+1:nP),SPD(m+1:nP),SPEV(m+1:nP))
    end if

end subroutine spavdmixed

subroutine spafreq(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    complex(C_DOUBLE_COMPLEX), allocatable :: af(:), raf(:)
    real(C_DOUBLE), allocatable :: a0(:), ra(:), w(:), w0(:)
    integer :: i, j, Nfft
    complex(C_DOUBLE_COMPLEX) :: IONE=cmplx(0.d0,1.d0)
    real(C_DOUBLE) :: w0i2, w0iwj, wj2
    type(C_PTR) :: plan, iplan

    SPA = 1.D0
    SPI = 1

    Nfft = nextpow2(n)

    allocate(a0(Nfft))
    allocate(af(Nfft))
    allocate(raf(Nfft))
    allocate(ra(Nfft))
    allocate(w(Nfft))
    allocate(w0(nP))

    call fftfreqs(Nfft,1.d0/dt,w)
    w = w*pi2
    w0 = pi2/P

    ra = 0.d0
    a0 = 0.d0
    a0(1:n) = acc

    plan  = fftw_plan_dft_r2c_1d(Nfft, a0, af, FFTW_ESTIMATE)
    iplan = fftw_plan_dft_c2r_1d(Nfft, af, ra, FFTW_ESTIMATE)
    call fftw_execute_dft_r2c(plan, a0, af)

    !$OMP PARALLEL DO PRIVATE(i, j, w0i2, w0iwj, wj2, raf, ra), SHARED(nP,Nfft,IONE,iplan,SPI,SPA)

    do i = 1, nP, 1
        w0i2 = w0(i)*w0(i)
        do j = 1, Nfft/2+1, 1
            w0iwj = w0(i)*w(j)
            wj2 = w(j)*w(j)
            raf(j) = af(j)*(w0i2+2.d0*zeta*w0iwj*IONE)/&
                    (w0i2-wj2+2.d0*zeta*w0iwj*IONE)
        end do
        call fftw_execute_dft_c2r(iplan, raf, ra)

        SPI(i) = maxloc(dabs(ra(1:n)),1)
        SPA(i) = ra(SPI(i))/dble(Nfft)
    end do

    !$OMP END PARALLEL DO
	
    call fftw_destroy_plan(plan)
    call fftw_destroy_plan(iplan)

    deallocate( a0)
    deallocate( af)
    deallocate(raf)
    deallocate( ra)
    deallocate(  w)
    deallocate( w0)

    return

end subroutine spafreq

subroutine pspafreq(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    complex(C_DOUBLE_COMPLEX), allocatable :: af(:), raf(:)
    real(C_DOUBLE), allocatable :: a0(:), ra(:), w(:), w0(:)
    integer :: i, j, Nfft
    complex(C_DOUBLE_COMPLEX) :: IONE=cmplx(0.d0,1.d0)
    real(C_DOUBLE) :: w0i2, w0iwj, wj2
    type(C_PTR) :: plan, iplan

    SPA = 1.D0
    SPI = 1

    Nfft = nextpow2(n)

    allocate(a0(Nfft))
    allocate(af(Nfft))
    allocate(raf(Nfft))
    allocate(ra(Nfft))
    allocate(w(Nfft))
    allocate(w0(nP))

    call fftfreqs(Nfft,1.d0/dt,w)
    w = w*pi2
    w0 = pi2/P

    ra = 0.d0
    a0 = 0.d0
    a0(1:n) = acc

    plan  = fftw_plan_dft_r2c_1d(Nfft, a0, af, FFTW_ESTIMATE)
    iplan = fftw_plan_dft_c2r_1d(Nfft, af, ra, FFTW_ESTIMATE)
    call fftw_execute_dft_r2c(plan, a0, af)

    do i = 1, nP, 1
        w0i2 = w0(i)*w0(i)
        do j = 1, Nfft/2+1, 1
            w0iwj = w0(i)*w(j)
            wj2 = w(j)*w(j)
            raf(j) = af(j)*(-1.d0/(w0i2-wj2+2.d0*zeta*w0iwj*IONE))
        end do
        call fftw_execute_dft_c2r(iplan, raf, ra)

        SPI(i) = maxloc(dabs(ra(1:n)),1)
        SPA(i) = ra(SPI(i))/dble(Nfft)*39.47841760435743D0/P(i)/P(i)
    end do

    call fftw_destroy_plan(plan)
    call fftw_destroy_plan(iplan)

    deallocate( a0)
    deallocate( af)
    deallocate(raf)
    deallocate( ra)
    deallocate(  w)
    deallocate( w0)

    return

end subroutine pspafreq

subroutine spavdfreq(acc,n,dt,zeta,P,nP,SPA,SPI,SPV,SPD,SPEV) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n), SPV(n), SPD(n), SPEV(n)
    integer, intent(out) :: SPI(n)

    complex(C_DOUBLE_COMPLEX), allocatable :: af(:), raf(:), rvf(:), rdf(:)
    real(C_DOUBLE), allocatable :: a0(:), ra(:), w(:), w0(:)
    real(C_DOUBLE), allocatable :: rv(:), rd(:)
    integer :: i, j, Nfft
    complex(C_DOUBLE_COMPLEX) :: IONE=cmplx(0.d0,1.d0)
    real(C_DOUBLE) :: w0i2, w0iwj, wj2, iNfft
    type(C_PTR) :: plan, iplan

    SPA = 1.D0
    SPV = 1.D0
    SPD = 1.D0
    SPI = 1

    Nfft = nextpow2(n)
    iNfft = 1.d0/dble(Nfft)

    allocate(a0(Nfft))
    allocate(af(Nfft))
    allocate(raf(Nfft))
    allocate(rvf(Nfft))
    allocate(rdf(Nfft))
    allocate(ra(Nfft))
    allocate(rv(Nfft))
    allocate(rd(Nfft))
    allocate(w(Nfft))
    allocate(w0(nP))

    call fftfreqs(Nfft,1.d0/dt,w)
    w = w*pi2
    w0 = pi2/P

    ra = 0.d0
    rv = 0.d0
    rd = 0.d0
    a0 = 0.d0
    a0(1:n) = acc

    plan  = fftw_plan_dft_r2c_1d(Nfft, a0, af, FFTW_ESTIMATE)
    iplan = fftw_plan_dft_c2r_1d(Nfft, af, ra, FFTW_ESTIMATE)
    call fftw_execute_dft_r2c(plan, a0, af)

    do i = 1, nP, 1
        w0i2 = w0(i)*w0(i)
        do j = 1, Nfft/2+1, 1
            w0iwj = w0(i)*w(j)
            wj2 = w(j)*w(j)
            raf(j) = af(j)*(w0i2+2.D0*zeta*w0iwj*IONE)/&
                    (w0i2-wj2+2.d0*zeta*w0iwj*IONE)
            rvf(j) = af(j)*(-w(j)*IONE/(w0i2-wj2+2.d0*zeta*w0iwj*IONE))
            rdf(j) = af(j)*(-1.d0/(w0i2-wj2+2.d0*zeta*w0iwj*IONE))
        end do
        call fftw_execute_dft_c2r(iplan, raf, ra)
        call fftw_execute_dft_c2r(iplan, rvf, rv)
        call fftw_execute_dft_c2r(iplan, rdf, rd)
        SPI(i) = maxloc(dabs(ra(1:n)),1)
        SPA(i) = dabs(ra(SPI(i)))*iNfft
        SPV(i) = maxval(dabs(rv(1:n)))*iNfft
        SPD(i) = maxval(dabs(rd(1:n)))*iNfft
        SPEV(i) = sqrt(2.d0*trapz(-rv(1:n)*acc*iNfft,n,dt,0.d0))
    end do

    call fftw_destroy_plan(plan)
    call fftw_destroy_plan(iplan)

    deallocate( a0)
    deallocate( af)
    deallocate(raf)
    deallocate(rvf)
    deallocate(rdf)
    deallocate( ra)
    deallocate( rv)
    deallocate( rd)
    deallocate(  w)
    deallocate( w0)

    return

end subroutine spavdfreq

subroutine spanmk(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    integer :: k
    real(8) :: SPVk, SPDk, SPEVk

    !$OMP PARALLEL DO PRIVATE(k)
    do k = 1, nP, 1
        call newmark(acc,n,dt,zeta,P(k),SPA(k),SPI(k),SPVk,SPDk,SPEVk)
    end do
    !$OMP END PARALLEL DO

end subroutine spanmk

subroutine pspanmk(acc,n,dt,zeta,P,nP,SPA,SPI) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(n)
    real(C_DOUBLE), intent(out) :: SPA(n)
    integer, intent(out) :: SPI(n)

    integer :: k
    real(8) :: SPVk, SPDk, SPEVk

    do k = 1, nP, 1
        call newmark(acc,n,dt,zeta,P(k),SPDk,SPI(k),SPVk,SPA(k),SPEVk)
        SPA(k) = SPA(k)*39.47841760435743D0/P(k)/P(k)
    end do

end subroutine pspanmk

subroutine spavdnmk(acc,n,dt,zeta,P,nP,SPA,SPI,SPV,SPD,SPEV) bind(c)

    integer, intent(in) :: n, nP
    real(C_DOUBLE), intent(in) :: dt, zeta
    real(C_DOUBLE), intent(in) :: acc(n), P(nP)
    real(C_DOUBLE), intent(out) :: SPA(nP), SPV(nP), SPD(nP), SPEV(nP)
    integer, intent(out) :: SPI(nP)

    integer :: k

    do k = 1, nP, 1
        call newmark(acc,n,dt,zeta,P(k),SPA(k),SPI(k),SPV(k),SPD(k),SPEV(k))
    end do

end subroutine spavdnmk

subroutine fitspectra(acc,n,dt,zeta,P,nP,SPAT,a,tol,mit) bind(c)
!     频域谱拟合
    integer, intent(in) :: n, nP, mit
    real(C_DOUBLE), intent(in) :: dt, zeta, tol
    real(C_DOUBLE), intent(in) :: acc(n), P(n), SPAT(:)
    real(C_DOUBLE), intent(out) :: a(:)

    real(C_DOUBLE), allocatable :: SPA(:), R(:), Rf(:), Pf(:)
    integer(C_INT), allocatable :: SPI(:)

    real(C_DOUBLE) :: aerror, merror, peak0, peak
    complex(C_DOUBLE_COMPLEX), allocatable :: af(:), afs(:)
    real(C_DOUBLE), allocatable :: a0(:), f(:)
    integer :: i, j, Nfft, NPf, IPf1, IPf2, iter, pk
    complex(C_DOUBLE_COMPLEX) :: IONE=cmplx(0.d0,1.d0)
    real(C_DOUBLE) :: pfi, iNfft, p0, pe0, dp, phi, t, ra
    type(C_PTR) :: plan, iplan

    peak0 = maxval(abs(acc), dim=1)
    Nfft = nextpow2(n)*2
    iNfft = 1.d0/dble(Nfft)

    allocate(a0(Nfft))
    allocate(af(Nfft))
    allocate(afs(Nfft))
    allocate(f(Nfft))
    allocate(Pf(Nfft/2))
    allocate(Rf(Nfft/2))
    a = acc

    call fftfreqs(Nfft,1.d0/dt,f)
    Pf(2:Nfft/2) = 1.d0/f(2:Nfft/2)
    Pf(1) = 2.d0*Pf(2)

    call decrfindfirst(Pf,Nfft/2,P(nP),IPf1)
    call decrfindlast(Pf,Nfft/2,P(1), IPf2)

    NPf = IPf2-IPf1+1

    a0 = 0.d0
    a0(1:n) = acc

    plan  = fftw_plan_dft_r2c_1d(Nfft, a0, af, FFTW_ESTIMATE)
    iplan = fftw_plan_dft_c2r_1d(Nfft, af, a0, FFTW_ESTIMATE)
    call fftw_execute_dft_r2c(plan, a0, af)

    allocate(SPA(nP))
    allocate(R(nP))
    allocate(SPI(nP))

    call spamixed(acc,n,dt,zeta,P,nP,SPA,SPI)
    call error(abs(SPA),SPAT,nP,aerror,merror)
    R = SPAT/abs(SPA)
    call decrlininterp(P,R,nP,Pf(IPf1:IPf2),Rf(IPf1:IPf2),NPf)
    write(unit=*, fmt="(A29,2F8.4)") "Initial Error: ",aerror,merror

    iter = 1
    do while ( (aerror>tol .or. merror>1.d0*tol) .and. iter<=mit )

        j = IPf2
        do i = 2, nP-1, 1
            p0 = p(i)
            t = dt*SPI(i)-dt
            dp = min(p0-P(i-1),P(i+1)-p0)

            do while ( Pf(j)<p0+0.5d0*dp .and. j>IPf1)
                pe0 = Pf(j)
                phi = atan2(aimag(af(j)),real(af(j)))
                ra = rsimple(pe0,p0,phi,t,zeta)
                if ( sign(1.d0,ra)*sign(1.d0,SPA(i)) < 0.d0) then
                    Rf(j) = 1.d0/Rf(j)
                end if
                j = j - 1
            end do
        end do

        af(IPf1:IPf2) = af(IPf1:IPf2)*Rf(IPf1:IPf2)
        afs = af
        call fftw_execute_dft_c2r(iplan, afs, a0)
        a = a0(1:n)*iNfft

        call adjustpeak(a,n,peak0)
        call adjustbaseline(a,n,dt)
        call adjustpeak(a,n,peak0)

        call spamixed(a,n,dt,zeta,P,nP,SPA,SPI)
        call error(abs(SPA),SPAT,nP,aerror,merror)
        R = SPAT/abs(SPA)
        call decrlininterp(P,R,nP,Pf(IPf1:IPf2),Rf(IPf1:IPf2),NPf)
        write(unit=*, fmt="(A11,I4,A14,2F8.4)") "Error After",iter,"  Iterations: ",aerror,merror
        iter = iter + 1
    end do

    call fftw_destroy_plan(plan)
    call fftw_destroy_plan(iplan)

    deallocate(SPA)
    deallocate(R)
    deallocate(Pf)
    deallocate(Rf)
    deallocate(SPI)
    deallocate(a0)
    deallocate(af)
    deallocate(afs)
    deallocate(f)

end subroutine fitspectra

subroutine adjustpeak(a,n,peak0) bind(c)
!     谱拟合过程中用到的幅值调整子程序
    integer, intent(in) :: n
    real(C_DOUBLE), intent(inout) :: a(:)
    real(C_DOUBLE), intent(in) :: peak0

    integer :: pk = 0, i
    real(C_DOUBLE) :: peak = 0.d0

    do i = 1, n, 1
        if ( abs(a(i))>peak ) then
            peak = a(i)
            pk = i
        end if
        if ( a(i)>peak0 ) then
            a(i) = peak0
        elseif ( a(i)<-peak0 ) then
            a(i) = -peak0
        end if
    end do

    if ( peak<peak0 ) then
        a(pk) = sign(peak0,peak)
    end if

end subroutine adjustpeak

subroutine adjustbaseline(a,n,dt) bind(c)
!     谱拟合过程中用到的基线调整子程序
    
    integer, intent(in) :: n
    real(C_DOUBLE), intent(inout) :: a(:)
    real(C_DOUBLE), intent(in) :: dt

    real(C_DOUBLE), allocatable :: v(:), d(:)
    integer, allocatable :: tp(:)
    integer :: ntp = 4, pl = 2, ph, i

    ph = pl + ntp - 1

    allocate(v(n))
    allocate(d(n))
    allocate(tp(ntp))

    do i = 1, ntp, 1
        tp(i) = int(dble(i)/dble(ntp)*dble(n))
    end do

    call ratacc2vd(a,v,d,n,dt,0.d0,0.d0)
    call targetdc(a,d,n,tp,ntp,ph,pl,dt,0.d0,0.d0)

    deallocate(v)
    deallocate(d)
    deallocate(tp)

    return
end subroutine adjustbaseline

subroutine adjustspectra(acc,n,dt,zeta,P,nP,SPAT,a,tol,mit) bind(c)
!     时域谱拟合
    integer, intent(in) :: n, nP, mit
    real(C_DOUBLE), intent(in) :: dt, zeta, tol
    real(C_DOUBLE), intent(in) :: acc(n), P(n), SPAT(:)
    real(C_DOUBLE), intent(out) :: a(:)

    real(C_DOUBLE), allocatable :: SPA(:), dR(:), ra(:,:,:)
    real(C_DOUBLE), allocatable :: M(:,:), W(:,:)
    integer(C_INT), allocatable :: SPI(:), SPIp(:)

    real(C_DOUBLE) :: aerror, merror, peak0
    integer :: i,j,iter

    allocate(SPA(nP))
    allocate(dR(nP))
    allocate(SPI(nP))
    allocate(SPIp(nP))
    allocate(ra(n,nP,nP))
    allocate(M(nP,nP))
    allocate(W(n,nP))

    peak0 = maxval(abs(acc), dim=1)
    a = acc
    SPIp = -1

    iter = 1
    call spamixed(a,n,dt,zeta,P,nP,SPA,SPI)
    call errora(abs(SPA),SPAT,nP,aerror,merror)
    write(unit=*, fmt="(A29,2F8.4)") "Initial Error: ",aerror,merror
    do while ( (aerror>tol .or. merror>1.d0*tol) .and. iter<=mit )

        dR = SPA*(SPAT/abs(SPA)-1.d0)/SPAT

        do i = 1, nP, 1
            if ( SPI(i) /= SPIp(i) ) then
                call wfunc( n,dt,SPI(i),P(i),zeta,W(:,i) )
            end if
        end do

        do i = 1, nP, 1
            do j = 1, nP, 1
                if ( SPI(j) /= SPIp(j) ) then
                    call ramixed(W(:,j),n,dt,zeta,P(i),ra(:,i,j))
                end if
                M(i,j) = ra(SPI(i),i,j)/SPAT(i)
                if ( i /= j ) then
                    M(i,j) = M(i,j)*0.618
                end if
            end do
        end do

        call leastsqs(M,dR,nP,nP)

        do i = 1, nP, 1
            a = a + dR(i)*W(:,i)
        end do

        call adjustpeak(a,n,peak0)
        call adjustbaseline(a,n,dt)
        call adjustpeak(a,n,peak0)

        SPIp = SPI
        call spamixed(a,n,dt,zeta,P,nP,SPA,SPI)
        call errora(abs(SPA),SPAT,nP,aerror,merror)
        write(unit=*, fmt="(A11,I4,A14,2F8.4)") "Error After",iter,"  Iterations: ",aerror,merror
        iter = iter + 1

    end do

    deallocate(SPA)
    deallocate(dR)
    deallocate(SPI)
    deallocate(SPIp)
    deallocate(ra)
    deallocate(M)
    deallocate(W)

end subroutine adjustspectra

subroutine wfunc(n,dt,itm,P,zeta,wf) bind(c)
!     时域谱拟合所需小波函数
    integer, intent(in) :: n, itm
    real(8), intent(in) :: dt, P, zeta
    real(8), intent(out) :: wf(:)

    real(8) :: tm, w, f, tmp1, gamma, deltaT
    real(8), allocatable :: tmp2(:)
    integer :: i

    allocate(tmp2(n))

    tm = dble(itm-1)*dt
    w = pi2/P
    f = 1.d0/P

    tmp1 = sqrt(1.d0-zeta**2)
    gamma = 1.178d0*(f*tmp1)**(-0.93d0)
    deltaT = atan(tmp1/zeta)/(w*tmp1)

    tmp2 = [ ( ( dble(i-1)*dt-tm+deltaT ),i=1,n ) ]

    wf = cos(w*tmp1*tmp2)*exp(-(tmp2/gamma)**2)

    deallocate(tmp2)

    return
end subroutine wfunc

end module spectra
